%=====================================================================================================================
% Generate the IRF after a positive fertility shock for the Example of Section 8 of LANCIA, RUSSO, and WORRALL (2024) %
%=====================================================================================================================

clear
close all
clc
delete *.asv

cd ..\
cd ..\
cd Stored_file

global wgrid w_thr wmax wmax_shock betta delta gamma pi1 pi2 e1 e2 e1_shock e2_shock ey1 ey2 eo1 eo2 v1aut v2aut wmin wfb nw n n_shock

load Canonical_policy

cd ..\
cd Matlab_functions
w_thr=w1(1);

%=========================================================================
%                           Endowment process
%=========================================================================
n=1;
n_shock=(1.01)*n;

e1 = n*ey1+eo1; % total endowment as share of mass of old in state 1
e2 = n*ey2+eo2; % total endowment as share of mass of old in state 2

e1_shock = n_shock*ey1+eo1; % total endowment as share of mass of old in state 1
e2_shock = n_shock*ey2+eo2; % total endowment as share of mass of old in state 2

%=========================================================================
%                           First best 
%=========================================================================

cy1fb=(delta./(betta)).*(n.*ey1+eo1)./(1+(delta/betta).*n);
cy2fb=(delta./(betta)).*(n.*ey2+eo2)./(1+(delta/betta).*n);
co1fb=n.*ey1+eo1-n.*cy1fb;
co2fb=n.*ey2+eo2-n.*cy2fb;
Ecyfb=pi1.*cy1fb+pi2.*cy2fb;

CY1fb=(cy1fb*n)./(n.*ey1+eo1); % aggregate consumption as share of tot endowment
CY2fb=(cy2fb*n)./(n.*ey2+eo2); % aggregate consumption as share of tot endowment
ECYfb=pi1.*CY1fb+pi2.*CY2fb;

cy1fb_shock=((delta./(betta))./(1+(delta/betta).*n_shock)).*(n_shock.*ey1+eo1);
cy2fb_shock=((delta./(betta))./(1+(delta/betta).*n_shock)).*(n_shock.*ey2+eo2);
co1fb_shock=n_shock.*ey1+eo1-n_shock.*cy1fb_shock;
co2fb_shock=n_shock.*ey2+eo2-n_shock.*cy2fb_shock;
Ecyfb_shock=pi1.*cy1fb_shock+pi2.*cy2fb_shock;

CY1fb_shock=(cy1fb_shock*n_shock)./(n_shock.*ey1+eo1); % aggregate consumption as share of tot endowment
CY2fb_shock=(cy2fb_shock*n_shock)./(n_shock.*ey2+eo2); % aggregate consumption as share of tot endowment
ECYfb_shock=pi1.*CY1fb_shock+pi2.*CY2fb_shock;

dfb1=1-(1/ey1).*(delta/(delta+betta));
dfb2=1-(1/ey2).*(delta/(delta+betta));

wfb = pi1*utility(co1fb,gamma)+pi2*utility(co2fb,gamma); % expected promised utility in the first-best 

%=========================================================================
%                           Policy after shocks 
%=========================================================================

nw=length(V);

% autarky levels of utility
[v1aut,v2aut,wmin] = Autarky(gamma,betta,pi1,pi2,ey1,ey2,eo1,eo2);

% maximum level of utility
[wmax,cy1_st,cy2_st] = Maxstat_n(gamma,betta,pi1,pi2,ey1,ey2,eo1,eo2,v1aut,v2aut,1);
[wmax_shock,cy1_st_shock,cy2_st_shock] = Maxstat_n(gamma,betta,pi1,pi2,ey1,ey2,eo1,eo2,v1aut,v2aut,n_shock);

% determination of policies
Vtilde=V;
tollv=1e-10;
distmia=1;

cy1_old=cy1;
cy2_old=cy2;
w1_old=w1;
w2_old=w2;

while distmia > tollv
[cy1shock,cy2shock,w1shock,w2shock,Vshock] = Value_Iteration_demo(cy1_old,cy2_old,w1_old,w2_old,V,Vtilde);
distmia=max(max(abs(Vshock-Vtilde)))

Vtilde=Vshock;
cy1_old=cy1shock;
cy2_old=cy2shock;
w1_old=w1shock;
w2_old=w2shock;

end

figure('name','Policy before/after shock')
subplot(2,2,1)
plot(wgrid,cy1,'b'); hold on;
plot(wgrid,cy1shock,'b--');

subplot(2,2,2)
plot(wgrid,cy2,'r'); hold on;
plot(wgrid,cy2shock,'r--'); hold on;
plot(wfb,0.5,'*')

subplot(2,2,3)
plot(wgrid,w1,'b'); hold on;
plot(wgrid,w1shock,'b--'); hold on;
plot(wgrid,w2,'r'); hold on;
plot(wgrid,w2shock,'r--'); hold on;
plot(wgrid,wgrid,'k:'); hold on;
plot(wfb,wfb,'*'); hold on;

subplot(2,2,4)
plot(wgrid,V,'b'); hold on;
plot(wgrid(1:end-1),Vshock(1:end-1),'b--'); hold on;

% next period consumption
cy11=interp1(wgrid,cy1,w1,'spline');
cy12=interp1(wgrid,cy2,w1,'spline');
cy21=interp1(wgrid,cy1,w2,'spline');
cy22=interp1(wgrid,cy2,w2,'spline');

% next period consumption:
% note that consumption next period goes back to normal
cy11shock=interp1(wgrid,cy1,w1shock,'spline');
cy12shock=interp1(wgrid,cy2,w1shock,'spline');
cy21shock=interp1(wgrid,cy1,w2shock,'spline');
cy22shock=interp1(wgrid,cy2,w2shock,'spline');

% CHECK BINDINGNESS OF PROMISE KEEPING
% if negative shock the PK becomes binding for lower level of w
% if positive shock the PK becomes binding for higher level of w

w_tilde=pi1.*log(n.*ey1+eo1-n.*cy1)+pi2.*log(n.*ey2+eo2-n.*cy2);
w_tilde_shock=pi1.*log(n_shock.*ey1+eo1-n_shock.*cy1shock)+pi2.*log(n_shock.*ey2+eo2-n_shock.*cy2shock);

plot(wgrid,w_tilde,'b',wgrid,wgrid,'k:',w_thr,w_thr,'*'); hold on;
plot(wgrid,w_tilde_shock,'r');

%================================================================================
%%%%%%%%%%%%%%%%%%% GENERATE DEBT POLICY %%%%%%%%%%%%%%%%%%%%%
%================================================================================

%----------DEBT BEFORE SHOCKS--------%
% state-contingent outstanding bond
b1 = (ey1-cy1)./ey1;
b2 = (ey2-cy2)./ey2;

% new issued state-contingent bond
b11 = (ey1-cy11)./ey1;
b12 = (ey2-cy12)./ey2;
b21 = (ey1-cy21)./ey1;
b22 = (ey2-cy22)./ey2;

%----------DEBT AFTER SHOCKS--------%
% state-contingent outstanding bond
b1shock = (ey1-cy1shock)./ey1;
b2shock = (ey2-cy2shock)./ey2;

% new issued state-contingent bond
b11shock = (ey1-cy11shock)./ey1;
b12shock = (ey2-cy12shock)./ey2;
b21shock = (ey1-cy21shock)./ey1;
b22shock = (ey2-cy22shock)./ey2;

%================================================
% Ergodic distribution
%================================================
nb=1000;
wmin = w_thr;
wmax = fsolve(@(x)interp1(wgrid,b12,x,'spline')-interp1(wgrid,b1,x,'spline'),-0.6); % it can makes errors  % it can generate errors!!
wmax_shock = fsolve(@(x)interp1(wgrid,b12shock,x,'spline')-interp1(wgrid,b1shock,x,'spline'),-0.6); % it can makes errors  % it can generate errors!!

wgridmod = [wmin:(max(wmax,wmax_shock)-wmin)./(nb-1):max(wmax,wmax_shock)];

for i=1:nb
debt(i)=interp1(wgrid,b1,wgridmod(i),'spline');
debt1(i)=interp1(wgrid,b11,wgridmod(i),'spline');
debt2(i)=interp1(wgrid,b12,wgridmod(i),'spline'); 
end

for i=1:nb
debtshock(i)=interp1(wgrid,b1shock,wgridmod(i),'spline');
debt1shock(i)=interp1(wgrid,b11shock,wgridmod(i),'spline');
debt2shock(i)=interp1(wgrid,b12shock,wgridmod(i),'spline'); 
end

figure('name','debt before and after shock in modified state space')
plot(debt,debt1,'b',debt,debt2,'r',debt,debt,'k:'); hold on
plot(debtshock,debt1shock,'b--',debtshock,debt2shock,'r--',debtshock,debtshock,'k:')

% discretize on the largest set
%DEBT=[min(min(debt),min(debtshock)):(max(max(debt),max(debtshock))-min(min(debt),min(debtshock)))/(nb-1):max(max(debt),max(debtshock))];
%edge= DEBT(1:nb); edge(nb+1)=max(DEBT);
edge= debt(1:nb); edge(nb+1)=max(debt);
debtk = discretize(debt,edge);
debt1k = discretize(debt1,edge); debt2k = discretize(debt2,edge);

debt1k(1)=1;
% correct
for i=2:nb
if isnan(debt1k(i))
   debt1k(i)=debt1k(i-1); 
end
if isnan(debt2k(i))
   debt2k(i)=debt2k(i-1);  
end
end

plot(debtk,debt1k,'b',debtk,debt2k,'r')

T11=zeros(nb,nb); T12=zeros(nb,nb); T21=zeros(nb,nb); T22=zeros(nb,nb);

% matrix of transition probabilities and state prices
for i=1:nb
   for j=1:nb
       if debtk(j)==debt1k(i)
            T11(i,j)=pi1; T21(i,j)=pi1;
       else
            T11(i,j)=0; T21(i,j)=0;
       end
       if debtk(j)==debt2k(i)
           T12(i,j)=1-pi1; T22(i,j)=1-pi1;
       else
           T12(i,j)=0; T22(i,j)=0; 
       end   
   end
end

T=[T11 T12; T21 T22]; T=T';

max(sum(T))

% stationary distribution
probb = (1/(2*nb))*ones(2*nb,1);
   test=1;
   while test > 10^(-10);
      probbx = T*probb; test = max(abs(probbx-probb));
      probb = probbx;
   end;

%probb_norm=probb/sum(probb); 

probb_norm=probb;

for i=1:nb
  probb1(i)=probb_norm(i);
end
sum(probb1)

for i=nb+1:2*nb
   probb2(i-nb)=probb_norm(i);
end
sum(probb2)

figure('name','long run distribution')
plot(debt1,probb1,'b',debt2,probb2,'r')
plot(debt,probb1+probb2,'k')

Esdebt(1)=probb1*(ey1.*debt)'+probb2*(ey2.*debt)';

% Adjustment of distribution after shock
debtkshock = discretize(debtshock,edge);

debtkshock(1)=1;

% correct
for i=2:nb    
if isnan(debtkshock(i))
   debtkshock(i)=debtkshock(i-1); 
end
end


prob1_shock=zeros(1,length(debtk));
prob2_shock=zeros(1,length(debtk));

% generate prob_shock
for i=1:length(debtk)
for j=1:length(debtk)
   if debtk(i)==debtkshock(j) && probb1(j)>0
       prob1_shock(i)=probb1(j);
   end
end
end

for i=1:length(debtk)
for j=1:length(debtk)
   if debtk(i)==debtkshock(j) && probb2(j)>0
       prob2_shock(i)=probb2(j);
   end
end
end

figure('name','Shocked distribution')
subplot(2,2,1)
plot(debtk,probb1,'b'); hold on
plot(debtkshock,probb1,'r'); hold on
plot(debtk,prob1_shock,'g')

subplot(2,2,2)
plot(debtk,probb2,'b'); hold on
plot(debtkshock,probb2,'r'); hold on
plot(debtk,prob2_shock,'g')

Esdebt(2)=(prob1_shock*(ey1.*debtshock)'+prob2_shock*(ey2.*debtshock)').*n_shock;

prob_shock=[prob1_shock prob2_shock];

% iteration until convergence
D1=[ey1.*(1-debt) ey2.*(1-debt)]*probb;

D1n=([ey1.*(1-debt).*n./(n.*ey1+eo1) ey2.*(1-debt).*n./(n.*ey2+eo2)]*probb);

D2=([ey1.*(1-debt) ey2.*(1-debt)]*prob_shock');
D2n=([ey1.*(1-debt).*n_shock./(n_shock.*ey1+eo1) ey2.*(1-debt).*n_shock./(n_shock.*ey2+eo2)]*prob_shock');

probb_post1 = T*prob_shock';
sum(probb_post1)
D3=[ey1.*(1-debt) ey2.*(1-debt)]*probb_post1;
D3n=([ey1.*(1-debt).*n./(n.*ey1+eo1) ey2.*(1-debt).*n./(n.*ey2+eo2)]*probb_post1);


probb_post2 = T*probb_post1;
sum(probb_post2)
D4=[ey1.*(1-debt) ey2.*(1-debt)]*probb_post2;
D4n=([ey1.*(1-debt).*n./(n.*ey1+eo1) ey2.*(1-debt).*n./(n.*ey2+eo2)]*probb_post2);

probb_post3 = T*probb_post2;
sum(probb_post3)
D5=[ey1.*(1-debt) ey2.*(1-debt)]*probb_post3;
D5n=([ey1.*(1-debt).*n./(n.*ey1+eo1) ey2.*(1-debt).*n./(n.*ey2+eo2)]*probb_post3);

probb_post4 = T*probb_post3;
sum(probb_post4)
D6=[ey1.*(1-debt) ey2.*(1-debt)]*probb_post4;
D6n=([ey1.*(1-debt).*n./(n.*ey1+eo1) ey2.*(1-debt).*n./(n.*ey2+eo2)]*probb_post4);

probb_post5 = T*probb_post4;
sum(probb_post5)
D7=[ey1.*(1-debt) ey2.*(1-debt)]*probb_post5;
D7n=([ey1.*(1-debt).*n./(n.*ey1+eo1) ey2.*(1-debt).*n./(n.*ey2+eo2)]*probb_post5);

probb_post6 = T*probb_post5;
sum(probb_post6)
D8=[ey1.*(1-debt) ey2.*(1-debt)]*probb_post6;
D8n=([ey1.*(1-debt).*n./(n.*ey1+eo1) ey2.*(1-debt).*n./(n.*ey2+eo2)]*probb_post6);

DEBTIRF=[D1 D2 D3 D4 D5 D6 D7 D8];
DEBTIRFn=[D1n D2n D3n D4n D5n D6n D7n D8n];

cd ..\
cd Stored_file
save('IRF_demo.mat','DEBTIRF','DEBTIRFn','Ecyfb','ECYfb', 'Ecyfb_shock', 'ECYfb_shock')

